package com.nvwa.dto.tree;

import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 树形结构相关方法，包括平铺和组装
 */
public final class TreeNodeUtils {
    /**
     * 默认顶级节点的父ID
     */
    static int TOP_NODE_PARENT_ID = 0;

    /**
     * 将平铺的节点转为树形结构
     *
     * @param topNodeParentId 顶级节点的父ID
     * @param nodeList        需要转换的节点集合
     * @param <ID>            节点唯一标识的类型
     * @param <T>             节点的类型
     * @return 树形结构的节点信息
     */
    public static <ID extends Number, T extends TreeNode<ID, T>> List<T> tree(ID topNodeParentId, List<T> nodeList) {
        if (nodeList == null || nodeList.isEmpty()) {
            return new ArrayList<>(0);
        }

        // 以父节点ID为key进行分组
        Map<ID, List<T>> groupingByParentIdMap = nodeList.stream()
                .collect(Collectors.groupingBy(TreeNode::getNodeParentId));

        // 获取顶级节点
        List<T> topNodeList = groupingByParentIdMap.remove(topNodeParentId);
        if (topNodeList == null || topNodeList.isEmpty()) {
            return new ArrayList<>(0);
        }

        // 对顶级节点进行排序，然后设置每个节点的子节点
        if (topNodeList.get(0).getNodeOrder() != null) {
            topNodeList.sort(Comparator.comparingInt(TreeNode::getNodeOrder));
        }
        topNodeList.forEach(topNode -> setChildren(topNode, groupingByParentIdMap));

        return topNodeList;
    }

    /**
     * 将树形结构的节点平铺后返回
     *
     * @param childNodeIdInitValue 子级节点ID初始值生成规则，参数为父节点ID，如果为null，则不会设置节点ID
     * @param topNodeParentId      顶级节点的父ID
     * @param treeNodeList         需要平铺的树形节点集合
     * @param <ID>                 节点唯一标识的类型
     * @param <T>                  节点的类型
     * @return 节点平铺后的节点集合
     */
    public static <ID extends Number, T extends TreeNode<ID, T>> List<T> flat(Function<Number, Number> childNodeIdInitValue,
                                                                              Number topNodeParentId,
                                                                              List<T> treeNodeList) {
        // 如果没有需要平铺的节点，直接结束
        if (treeNodeList == null || treeNodeList.isEmpty()) {
            return new ArrayList<>(0);
        }

        // 子节点ID初始值
        Number initialCodeValue = Optional.ofNullable(childNodeIdInitValue)
                .map(initValue -> initValue.apply(topNodeParentId))
                .orElse(null);
        // 判断是否需要生成节点ID
        boolean generateNodeId = initialCodeValue != null;
        AtomicLong nodeId = new AtomicLong(0);
        if (generateNodeId) {
            // 设置节点ID初始值
            nodeId.set(initialCodeValue.longValue());
        }
        // 节点展示顺序
        AtomicInteger nodeOrder = new AtomicInteger(0);

        // 遍历树形节点，平铺每个节点的子节点
        List<T> children = treeNodeList.stream()
                .peek(treeNode -> {
                    if (generateNodeId) {
                        treeNode.setNodeId(nodeId.getAndIncrement());
                    }
                    treeNode.setNodeParentId(topNodeParentId);
                    treeNode.setNodeOrder(nodeOrder.getAndIncrement());
                })
                .map(treeNode -> flat(childNodeIdInitValue, treeNode.getNodeId(), treeNode.getNodeChildren()))
                .flatMap(Collection::stream)
                .collect(Collectors.toList());

        // 将平铺后的子节点添加当前节点集合中
        treeNodeList.addAll(children);

        return treeNodeList;
    }

    /**
     * 将树形结构的节点平铺后返回，子级节点ID生成规则为parentNodeId * 100 + 10递增
     *
     * @param topNodeParentId 顶级节点的父ID
     * @param treeNodeList    需要平铺的树形节点集合
     * @param <ID>            节点唯一标识的类型
     * @param <T>             节点的类型
     * @return 节点平铺后的节点集合
     */
    public static <ID extends Number, T extends TreeNode<ID, T>> List<T> flat(Number topNodeParentId, List<T> treeNodeList) {
        return flat(parentNodeId -> parentNodeId.longValue() * 100 + 10, topNodeParentId, treeNodeList);
    }

    /**
     * 将树形结构的节点平铺后返回，子级节点ID生成规则为parentNodeId * 100 + 10递增
     *
     * @param treeNodeList 需要平铺的树形节点集合
     * @param <ID>         节点唯一标识的类型
     * @param <T>          节点的类型
     * @return 节点平铺后的节点集合
     */
    public static <ID extends Number, T extends TreeNode<ID, T>> List<T> flat(List<T> treeNodeList) {
        return flat(TOP_NODE_PARENT_ID, treeNodeList);
    }

    /**
     * 设置指定节点的子节点
     *
     * @param node                  需要设置子节点的节点
     * @param groupingByParentIdMap 以父节点ID为key进行分组的Map
     * @param <ID>                  节点唯一标识的类型
     * @param <T>                   节点的类型
     */
    private static <ID extends Number, T extends TreeNode<ID, T>> void setChildren(T node, Map<ID, List<T>> groupingByParentIdMap) {
        // 获取当前节点的子节点
        List<T> children = groupingByParentIdMap.get(node.getNodeId());

        // 如果当前节点没有子节点，结束递归
        if (children == null || children.isEmpty()) {
            return;
        }

        // 对当前节点进行排序
        if (node.getNodeOrder() != null) {
            children.sort(Comparator.comparingInt(TreeNode::getNodeOrder));
        }
        // 设置当前节点的子节点
        node.setNodeChildren(children);

        // 设置当前节点子节点的子节点
        children.forEach(parentNode -> setChildren(parentNode, groupingByParentIdMap));
    }
}
